import os
import cv2
import json
import numpy as np
from paddleocr import PaddleOCR
# import layoutparser as lp  # 暂时移除，使用简化版布局分析

class ImageOCRProcessor:
    def __init__(self, lang='ch'):
        # 初始化OCR引擎
        self.ocr = PaddleOCR(use_angle_cls=True, lang=lang)
        
        # 使用简化版布局分析
        print("使用简化版布局分析方法")
    
    def process_image(self, image_path, output_dir=None):
        """处理图片文件，提取文本和布局信息"""
        if output_dir is None:
            output_dir = os.path.dirname(image_path)
        os.makedirs(output_dir, exist_ok=True)
        
        # 处理图像
        result = self.process_single_image(image_path)
        
        # 保存结果为JSON
        base_name = os.path.basename(image_path).split('.')[0]
        output_json = os.path.join(output_dir, f"{base_name}_result.json")
        with open(output_json, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
            
        return output_json
    
    def process_single_image(self, img_path):
        """处理单个图像，提取文本和布局"""
        # 读取图像
        image = cv2.imread(img_path)
        if image is None:
            raise ValueError(f"无法读取图像: {img_path}")
        
        # 预处理图像（可选）
        # image = self.preprocess_image(image)
        
        # OCR识别
        ocr_result = self.ocr.ocr(img_path, cls=True)
        
        # 简化版布局分析 - 使用基于OCR结果的启发式方法
        height, width = image.shape[:2]
        
        # 构建结果
        result = {
            "image": os.path.basename(img_path),
            "blocks": []
        }
        
        # 简化版布局分析 - 基于OCR结果进行分组
        result["blocks"] = []
        
        if ocr_result[0]:
            # 所有文本框
            all_texts = []
            for line in ocr_result[0]:
                bbox = line[0]
                text = line[1][0]
                confidence = line[1][1]
                all_texts.append({
                    "text": text,
                    "confidence": float(confidence),
                    "bbox": [[float(p[0]), float(p[1])] for p in bbox]
                })
            
            # 尝试检测表格区域 (简化版)
            # 这里使用启发式方法：检查文本框的排列是否规则
            table_candidates = self.detect_tables_heuristic(all_texts, width, height)
            
            # 添加检测到的表格
            for table in table_candidates:
                table_texts = []
                for text in all_texts:
                    # 检查文本是否在表格区域内
                    text_x1 = text["bbox"][0][0]
                    text_y1 = text["bbox"][0][1]
                    if (table[0] <= text_x1 <= table[2] and table[1] <= text_y1 <= table[3]):
                        table_texts.append(text)
                
                if table_texts:
                    table_content = self.extract_table_structure(
                        image[int(table[1]):int(table[3]), int(table[0]):int(table[2])], 
                        table_texts
                    )
                    result["blocks"].append({
                        "type": "table",
                        "coordinates": [int(coord) for coord in table],
                        "content": table_content
                    })
            
            # 剩余文本作为普通文本块
            remaining_texts = []
            for text in all_texts:
                is_in_table = False
                for table in table_candidates:
                    text_x1 = text["bbox"][0][0]
                    text_y1 = text["bbox"][0][1]
                    if (table[0] <= text_x1 <= table[2] and table[1] <= text_y1 <= table[3]):
                        is_in_table = True
                        break
                
                if not is_in_table:
                    remaining_texts.append(text)
            
            if remaining_texts:
                # 按Y坐标分组为段落
                paragraphs = self.group_texts_into_paragraphs(remaining_texts)
                
                for i, para_texts in enumerate(paragraphs):
                    # 计算段落边界
                    min_x = min(text["bbox"][0][0] for text in para_texts)
                    min_y = min(text["bbox"][0][1] for text in para_texts)
                    max_x = max(text["bbox"][2][0] for text in para_texts)
                    max_y = max(text["bbox"][2][1] for text in para_texts)
                    
                    # 判断是否为标题 (简单启发式：短且位于顶部)
                    is_title = (len(para_texts) <= 2 and min_y < height * 0.2)
                    
                    result["blocks"].append({
                        "type": "title" if is_title else "text",
                        "coordinates": [int(min_x), int(min_y), int(max_x), int(max_y)],
                        "content": para_texts
                    })
        
        # 如果布局分析没有找到任何块，则直接使用OCR结果
        if not result["blocks"] and ocr_result[0]:
            all_texts = []
            for line in ocr_result[0]:
                text = line[1][0]
                confidence = line[1][1]
                bbox = line[0]
                all_texts.append({
                    "text": text,
                    "confidence": float(confidence),
                    "bbox": [[float(p[0]), float(p[1])] for p in bbox]
                })
            
            result["blocks"].append({
                "type": "text",
                "coordinates": [0, 0, image.shape[1], image.shape[0]],
                "content": all_texts
            })
        
        return result
    
    def preprocess_image(self, image):
        """图像预处理，提高OCR质量"""
        # 转为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 二值化
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        
        # 降噪（可选）
        denoised = cv2.fastNlMeansDenoising(binary, None, 10, 7, 21)
        
        return denoised
    
    def extract_table_structure(self, table_image, texts):
        """尝试提取表格结构"""
        # 这里可以实现更复杂的表格结构提取算法
        # 简单实现：根据文本Y坐标分组为行
        if not texts:
            return []
            
        # 按Y坐标排序
        texts.sort(key=lambda x: x["bbox"][0][1])
        
        # 简单的行分组（可以改进）
        rows = []
        current_row = [texts[0]]
        row_y = texts[0]["bbox"][0][1]
        
        for text in texts[1:]:
            text_y = text["bbox"][0][1]
            # 如果Y坐标相差不大，认为是同一行
            if abs(text_y - row_y) < 20:  # 阈值可调整
                current_row.append(text)
            else:
                # 按X坐标排序当前行
                current_row.sort(key=lambda x: x["bbox"][0][0])
                rows.append([t["text"] for t in current_row])
                # 开始新行
                current_row = [text]
                row_y = text_y
        
        # 处理最后一行
        if current_row:
            current_row.sort(key=lambda x: x["bbox"][0][0])
            rows.append([t["text"] for t in current_row])
            
        return rows

    def visualize_result(self, image_path, result, output_path=None):
        """可视化OCR和布局分析结果"""
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图像: {image_path}")
        
        # 绘制布局块
        for block in result["blocks"]:
            x1, y1, x2, y2 = block["coordinates"]
            
            # 根据块类型选择颜色
            if block["type"] == "table":
                color = (0, 0, 255)  # 红色表示表格
            elif block["type"] == "title":
                color = (255, 0, 0)  # 蓝色表示标题
            else:
                color = (0, 255, 0)  # 绿色表示文本
                
            # 绘制矩形
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            
            # 添加类型标签
            cv2.putText(image, block["type"], (x1, y1-10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
            
            # 绘制文本框（可选）
            if block["type"] != "table":
                for text_item in block["content"]:
                    if "bbox" in text_item:
                        bbox = text_item["bbox"]
                        pts = np.array(bbox, np.int32).reshape((-1, 1, 2))
                        cv2.polylines(image, [pts], True, (255, 255, 0), 1)
        
        # 保存结果
        if output_path is None:
            base_dir = os.path.dirname(image_path)
            base_name = os.path.basename(image_path).split('.')[0]
            output_path = os.path.join(base_dir, f"{base_name}_visualized.jpg")
            
        cv2.imwrite(output_path, image)
        return output_path


    def detect_tables_heuristic(self, texts, width, height):
        """使用启发式方法检测表格
        
        基本思路：
        1. 寻找排列整齐的文本行
        2. 检查这些行是否有相似的宽度和间距
        """
        if not texts:
            return []
            
        # 按Y坐标排序
        texts_by_y = sorted(texts, key=lambda x: x["bbox"][0][1])
        
        # 分组为行
        rows = []
        current_row = [texts_by_y[0]]
        row_y = texts_by_y[0]["bbox"][0][1]
        
        for text in texts_by_y[1:]:
            text_y = text["bbox"][0][1]
            # 如果Y坐标相差不大，认为是同一行
            if abs(text_y - row_y) < 20:  # 阈值可调整
                current_row.append(text)
            else:
                # 按X坐标排序当前行
                current_row.sort(key=lambda x: x["bbox"][0][0])
                rows.append(current_row)
                # 开始新行
                current_row = [text]
                row_y = text_y
        
        # 处理最后一行
        if current_row:
            current_row.sort(key=lambda x: x["bbox"][0][0])
            rows.append(current_row)
        
        # 检测表格特征
        table_candidates = []
        
        # 至少需要3行才考虑是表格
        if len(rows) >= 3:
            # 检查行是否有相似的列数和间距
            row_lengths = [len(row) for row in rows]
            
            # 如果多行有相同的列数，可能是表格
            for i in range(len(rows) - 2):
                if row_lengths[i] >= 2 and row_lengths[i] == row_lengths[i+1] == row_lengths[i+2]:
                    # 找到可能的表格，计算边界
                    table_rows = rows[i:i+3]
                    min_x = min(text["bbox"][0][0] for row in table_rows for text in row)
                    min_y = min(text["bbox"][0][1] for row in table_rows for text in row)
                    max_x = max(text["bbox"][2][0] for row in table_rows for text in row)
                    max_y = max(text["bbox"][2][1] for row in table_rows for text in row)
                    
                    # 扩展表格边界，包含后续相似行
                    j = i + 3
                    while j < len(rows) and row_lengths[j] == row_lengths[i]:
                        for text in rows[j]:
                            min_x = min(min_x, text["bbox"][0][0])
                            min_y = min(min_y, text["bbox"][0][1])
                            max_x = max(max_x, text["bbox"][2][0])
                            max_y = max(max_y, text["bbox"][2][1])
                        j += 1
                    
                    # 添加边距
                    padding = 10
                    table_candidates.append([
                        max(0, min_x - padding),
                        max(0, min_y - padding),
                        min(width, max_x + padding),
                        min(height, max_y + padding)
                    ])
                    
                    # 跳过已处理的行
                    i = j - 1
        
        return table_candidates
    
    def group_texts_into_paragraphs(self, texts):
        """将文本分组为段落"""
        if not texts:
            return []
            
        # 按Y坐标排序
        texts_by_y = sorted(texts, key=lambda x: x["bbox"][0][1])
        
        # 分组为段落
        paragraphs = []
        current_para = [texts_by_y[0]]
        para_y = texts_by_y[0]["bbox"][2][1]  # 使用底部Y坐标
        
        for text in texts_by_y[1:]:
            text_y = text["bbox"][0][1]  # 使用顶部Y坐标
            # 如果新文本的顶部与当前段落底部距离较近，认为是同一段落
            if abs(text_y - para_y) < 30:  # 段落间距阈值
                current_para.append(text)
                # 更新段落底部坐标
                para_y = max(para_y, text["bbox"][2][1])
            else:
                # 按X坐标排序当前段落中的文本
                current_para.sort(key=lambda x: x["bbox"][0][0])
                paragraphs.append(current_para)
                # 开始新段落
                current_para = [text]
                para_y = text["bbox"][2][1]
        
        # 处理最后一个段落
        if current_para:
            current_para.sort(key=lambda x: x["bbox"][0][0])
            paragraphs.append(current_para)
            
        return paragraphs

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Image OCR and Layout Analysis')
    parser.add_argument('image_path', help='Path to the image file')
    parser.add_argument('--output', '-o', help='Output directory', default=None)
    parser.add_argument('--lang', '-l', help='Language for OCR', default='ch')
    parser.add_argument('--visualize', '-v', action='store_true', help='Visualize results')
    
    args = parser.parse_args()
    
    processor = ImageOCRProcessor(lang=args.lang)
    result_path = processor.process_image(args.image_path, args.output)
    
    print(f"Processing complete. Results saved to: {result_path}")
    
    if args.visualize:
        # 读取结果
        with open(result_path, 'r', encoding='utf-8') as f:
            result = json.load(f)
        
        # 可视化
        vis_path = processor.visualize_result(args.image_path, result)
        print(f"Visualization saved to: {vis_path}") 